# python-Levenshtein
from tqdm import tqdm

from utils_file import read_transcript, read_csv, get_csv_txt
from metrics import calculate_cer_en_zh, calculate_cer, calculate_wer
from constant import EOS_CHAR, SOS_CHAR, PAD_CHAR
from helper import get_word_segments_per_language

def calculate_correct(hyp_gold_dict):

    total_word, total_char, total_cer, total_wer = 0, 0, 0, 0
    total_en_cer, total_zh_cer, total_en_char, total_zh_char = 0, 0, 0, 0

    pbar = tqdm(iter(hyp_gold_dict.values()), leave=True, total=len(hyp_gold_dict))
    for i, data in enumerate(pbar):
        hyp, golden = data

        hyp = hyp.replace(EOS_CHAR, "").replace(SOS_CHAR, "").replace(PAD_CHAR, "")
        golden = golden.replace(EOS_CHAR, "").replace(SOS_CHAR, "").replace(PAD_CHAR, "")

        wer = calculate_wer(hyp, golden)
        cer = calculate_cer(hyp.strip(), golden.strip())

        en_cer, zh_cer, num_en_char, num_zh_char = calculate_cer_en_zh(hyp, golden)

        total_en_cer += en_cer
        total_zh_cer += zh_cer
        total_en_char += num_en_char
        total_zh_char += num_zh_char

        total_wer += wer
        total_cer += cer
        total_word += len(golden.split(" "))
        # total_char += len(golden)
        # print(golden)
        # total_word += len(get_word_segments_per_language(golden))
        total_char += len(get_word_segments_per_language(golden))



        pbar.set_description("TEST CER:{:.2f}% WER:{:.2f}% CER_EN:{:.2f}% CER_ZH:{:.2f}%".format(
                total_cer*100/total_char, total_wer*100/total_word, total_en_cer*100/max(1, total_en_char), total_zh_cer*100/max(1, total_zh_char)))

def main():

    file_dir = "/public/ai_platform/models/cmcc_2022/dataset_aishell/transcript"
    txt_dict = read_transcript(file_dir=file_dir)

    csv_dir = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/asr_conformer/tensorrt"
    csv_dir = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/conformer_onnxruntime/onnx/"
    csv_dir = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/conformer_onnxruntime/pytorch/"
    csv_dir = "/home/gyf/pkg/conformer_new/asr_conformer/speech_recognition/"
    # csv_dir = "/home/gyf/pkg/conformer_new/asr_conformer/tensorrt/"
    csv_file = "result_test.csv"
    # csv_file = "result_test_2023.10.22.csv"
   
    # csv_file = "result_test_pth.csv"
    # csv_file = "result_dev.csv"
    # csv_file = "/home/gyf/pkg/conformer_new/asr_conformer/tensorrt/result_test.csv"
    
    csv_dict = read_csv(csv_dir, csv_file)

    # print(csv_dict)
    # print(txt_dict)
    merge_dict = get_csv_txt(csv_dict, txt_dict)

    calculate_correct(merge_dict)

if __name__ == '__main__':
    main()